package com.maweirdos.ai.core.memory;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.maweirdos.ai.entity.ChatMessage;
import com.maweirdos.ai.mapper.ChatMessageMapper;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.stereotype.Repository;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.StringUtils;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 聊天记录数据库存储
 *
 * @author maweirdos
 * Create by 2025/7/2 11:34
 */
@Repository
@Slf4j
public class DatabaseChatMemoryRepository implements ChatMemoryRepository {

    private final ChatMessageMapper chatMessageMapper;

    public DatabaseChatMemoryRepository(ChatMessageMapper chatMessageMapper) {
        this.chatMessageMapper = chatMessageMapper;
    }

    @NotNull
    @Override
    public List<String> findConversationIds() {
        // 查询所有不同的会话ID
        return chatMessageMapper.selectList(null).stream()
                .map(ChatMessage::getConversationId)
                .distinct()
                .toList();
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        if (!StringUtils.hasText(conversationId)) {
            return List.of();
        }

        try {
            // 按创建时间排序查询指定会话ID的消息
            LambdaQueryWrapper<ChatMessage> queryWrapper = new LambdaQueryWrapper<>();
            queryWrapper.eq(ChatMessage::getConversationId, conversationId)
                    .orderByAsc(ChatMessage::getCreateTime);
            
            List<ChatMessage> chatMessages = chatMessageMapper.selectList(queryWrapper);
            
            // 将数据库实体转换为Spring AI消息对象
            List<Message> messages = new ArrayList<>();
            for (ChatMessage chatMessage : chatMessages) {
                // 直接根据消息类型和内容创建消息对象
                String messageType = chatMessage.getMessageType();
                String content = chatMessage.getContent();
                
                Message message = null;
                switch (messageType) {
                    case "user":
                        message = new UserMessage(content);
                        break;
                    case "assistant":
                        message = new AssistantMessage(content);
                        break;
                    case "system":
                        message = new SystemMessage(content);
                        break;
                }
                
                if (message != null) {
                    messages.add(message);
                }
            }
            
            return messages;
        } catch (Exception e) {
            log.error("根据会话ID查询消息异常: {}", conversationId, e);
            return List.of();
        }
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        if (!StringUtils.hasText(conversationId) || messages == null || messages.isEmpty()) {
            return;
        }

        try {
            for (Message message : messages) {
                // 获取消息内容
                String content = getMessageContent(message);

                // 将完整消息序列化为JSON
                Map<String, Object> rawMessageJson = new HashMap<>();
                rawMessageJson.put("type", message.getMessageType().getValue());
                rawMessageJson.put("content", content);
                rawMessageJson.put("metadata", message.getMetadata());

                // 将Spring AI消息对象转换为数据库实体
                ChatMessage chatMessage = ChatMessage.builder()
                        .conversationId(conversationId)
                        .userId("system") // 默认用户ID，实际应用中应该从上下文或参数中获取
                        .messageType(message.getMessageType().getValue())
                        .content(content)
                        .createTime(LocalDateTime.now())
                        .rawMessageJson(rawMessageJson)
                        .build();

                // 保存到数据库
                int result = chatMessageMapper.insert(chatMessage);
                log.info("保存消息结果: conversationId={}, rawMessageJson={}, result={}",
                        conversationId, rawMessageJson, result);
            }
        } catch (Exception e) {
            log.error("根据会话id保存消息异常: {}", conversationId, e);
            throw new RuntimeException("保存消息失败", e);
        }
    }

    @Override
    public void deleteByConversationId(String conversationId) {
        if (!StringUtils.hasText(conversationId)) {
            return;
        }

        try {
            LambdaQueryWrapper<ChatMessage> queryWrapper = new LambdaQueryWrapper<>();
            queryWrapper.eq(ChatMessage::getConversationId, conversationId);
            chatMessageMapper.delete(queryWrapper);
        } catch (Exception e) {
            log.error("根据会话ID删除消息异常: {}", conversationId, e);
        }
    }
    
    /**
     * 根据消息类型获取内容
     */
    private String getMessageContent(Message message) {
        if (message instanceof UserMessage) {
            return ((UserMessage) message).getText();
        } else if (message instanceof AssistantMessage) {
            return ((AssistantMessage) message).getText();
        } else if (message instanceof SystemMessage) {
            return ((SystemMessage) message).getText();
        }
        return "";
    }
}
